import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import DataLoader
import math
from sklearn.metrics import roc_curve, auc
import matplotlib.pyplot as plt
import pandas as pd
import os

from MyDataSet import MyDataSet

"""
CNN网络结构类型定义
尝试1次卷积
输入类型的数据是 71*47的矩阵，我认为不用池化层，尝试直接
"""
class MyNet(nn.Module):
    def __init__(self, label_lenth,em_num,ex_num,in_chanels = 1,out_chanels = 3, kernel_size = 4):
        super().__init__()
        self.conv1 = nn.Conv2d(in_chanels, out_chanels, kernel_size)
        self.pool1 = nn.MaxPool2d(2)
        self.out_chanels = out_chanels
        self.em_num = int((em_num-(kernel_size-1))/2)
        self.ex_num = int((ex_num-(kernel_size-1))/2)
        self.fc1 = nn.Linear(out_chanels * self.em_num * self.ex_num,100)
        self.fc2 = nn.Linear(100,50)
        self.fc3 = nn.Linear(50,label_lenth)

    # 前向传播
    def forward(self,x):
        x = torch.relu(self.conv1(x))
        x = self.pool1(x)
        x = x.view(-1,self.ex_num*self.em_num*self.out_chanels)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = F.log_softmax(self.fc3(x), dim = 1)
        return x

    def get_features(self,x):
        x = torch.relu(self.conv1(x))
        x = self.pool1(x)
        x = x.view(-1,self.ex_num*self.em_num*self.out_chanels)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return x

"""
定义网络得分类对象
便于取得分
"""
class NetGrades():
    def __init__(self,conf_martix,kind_count,acc,pre_list,rec_list,f1_list):
        self.conf_mar = conf_martix
        self.kind_count = kind_count
        self.acc = acc
        self.pre_list = pre_list
        self.rec_list = rec_list
        self.f1_list = f1_list

    def get_average(self):
        return self.acc,sum(self.pre_list)/self.kind_count,sum(self.rec_list)/self.kind_count,sum(self.pre_list)/self.kind_count,sum(self.f1_list)/self.kind_count

    def save_conf_mar(self,file_name,dir_root = r"D:\Code\CNN识别菌种\RunningData\ConfusionMartix"):
        df = pd.DataFrame(self.conf_mar)
        save_path = os.path.join(dir_root,file_name+".xlsx")
        df.to_excel(save_path)
"""
因为label是字符串类型，而网络输出值是一个数组类型，得把值都对应上，因此要建立一个映射字典
后期可能会更新菌种，所以要建立动态字典
int -> str的映射
"""
def creat_net_map(label_list):
    i = 0
    map = {}
    KindIndexMap = {}
    for label in label_list:
        if label not in KindIndexMap:
            map[i] = label
            KindIndexMap[label] = i
            i += 1
    return map,KindIndexMap


def get_data_loader(data_set,batch_size = 1):
    return DataLoader(data_set,batch_size=batch_size,shuffle=True)

def evaluate(test_data,net,net_map): #评估,打印版
    print(net_map)
    n_correct = 0
    n_total = 0

    with torch.no_grad():
        for (x,y) in test_data:
            outputs = net.forward(x)
            for i,output in enumerate(outputs):
                predict = net_map[int(torch.argmax(output))]
                real = y[i]
                if predict == real:
                    n_correct += 1
                n_total += 1
                print(output)
                print("Predict :  ",predict,"Real kind: ",real,"---Result of predict:",predict == real)
    print("Predict Count：", n_total,"Right Count：",n_correct,"Right Accuracy: ",n_correct/n_total)
    return n_correct/n_total

#返回混淆矩阵
def get_conf_martix(test_data,net,KindIndexMap):
    kind_count = len(KindIndexMap)
    ret = np.zeros((kind_count,kind_count))
    with torch.no_grad():
        for (x,y) in test_data:
            outputs = net.forward(x)
            for i,output in enumerate(outputs):
                predict = int(torch.argmax(output))
                real = KindIndexMap[y[i]]
                ret[real][predict] = ret[real][predict]+1
    return ret

#返回准确率，精确率(一维矩阵)，召回率和F1分数
def get_grades(test_data,net,KindIndexMap):
    conf_mar = get_conf_martix(test_data,net,KindIndexMap)
    total = np.sum(conf_mar)
    sum_row = np.sum(conf_mar,axis=1) #沿着行求和
    sum_col = np.sum(conf_mar,axis=0) #沿着列求和
    kind_count = len(KindIndexMap)

    #求准确率
    acc = 0
    for i in range(kind_count):
        acc = acc + conf_mar[i][i]/total

    #求精确率
    pre_list = []
    for i in range(kind_count): #求第i个物种的精确率
        if sum_col[i] == 0:
            pre_list.append(-1)
        else:
            pre_list.append(conf_mar[i][i]/sum_col[i])

    #求召回率
    rec_list = []
    for i in range(kind_count): #求第i个物种的召回率
        rec_list.append(conf_mar[i][i]/sum_row[i])

    #求F1分数
    f1_list = []
    for i in range(kind_count): #求第i个物种的F1分数
        if rec_list[i] == 0 or pre_list[i] == 0 or pre_list[i] == -1:
            f1_list.append(-1)
        else:
            f1_list.append(2/((1/rec_list[i])+(1/pre_list[i])))

    return NetGrades(conf_mar,kind_count,acc,pre_list,rec_list,f1_list)

#画出模型的ROC曲线，输入参数数据集和训练好的网络还有菌种名 - 数字映射图
def draw_ROC(test_data,net,KindIndexMap,net_map,file_name):
    kind_count = len(KindIndexMap)
    labels = []
    scores = []
    with torch.no_grad():
        for (x,y) in test_data:
            outputs = net.forward(x)
            for i,output in enumerate(outputs):
                labels.append(KindIndexMap[y[i]])
                scores.append([])
                # 计算得分的output，将数据转化为概率
                for val in output:
                    scores[-1].append(math.exp(val))
    labels = np.array(labels)
    scores = np.array(scores)

    #开始画每个类别的ROC曲线
    fprs = dict()
    tprs = dict()
    roc_auc = dict()

    for kind_key in range(kind_count): #画每个物种的曲线
        kind = net_map[kind_key]
        fprs[kind],tprs[kind],_ = roc_curve(labels==kind_key,scores[:,kind_key])
        roc_auc[kind] = auc(fprs[kind],tprs[kind])

    plt.figure(figsize=(12.8,7.2))
    for kind_key in range(kind_count):
        kind = net_map[kind_key]
        plt.plot(fprs[kind],tprs[kind],label = "ROC curve of class "+kind)

    plt.plot([0,1],[0,1],"k--")
    plt.xticks(fontsize=20)  # x轴刻度标签字体大小
    plt.yticks(fontsize=20)  # y轴刻度标签字体大小
    plt.xlabel('False Positive Rate',fontdict={"family": "Times New Roman"}, fontsize=32)
    plt.ylabel('True Positive Rate',fontdict={"family": "Times New Roman"}, fontsize=32)
    plt.title('ROC Curves for Multi-class Data',fontdict={"family": "Times New Roman"},fontsize=36)
    plt.legend(loc="lower right", fontsize=16)
    plt.savefig(r"D:\Code\CNN识别菌种\RunningData\ROC\\"+file_name+".png")

    return roc_auc









